Description

Implements Jupyter AI IPython magics via LangChain. The most common usage syntax looks like this:

%%ai <model-id> --format <format>
<prompt>

Usage

Rapid usage with OpenAI (<5 mins)

First, install the extension.

pip install jupyter_ai

Then, get an OpenAI API key. You can now use the magics in JupyterLab by starting the server with the OPENAI_API_KEY environment variable:

OPENAI_API_KEY=<key> jupyter lab

Then, in your notebook, load the extension (you must repeat this every time you start a new IPython session):

%load_ext jupyter_ai

You can now interact with models via the %%ai cell magic.

%%ai chatgpt
Write me a numerical algorithm that computes pi iteratively.

We support maintaining conversation transcripts out-of-the-box for OpenAI Chat models, meaning you can also ask it questions about previous exchanges:

%%ai chatgpt
Can you elaborate more on how this numerical method works?

Installation

Obviously, you'll need to install this package first:

pip install jupyter_ai

LangChain does not list every model provider as an explicit dependency, which means you may have to install them manually before using them. For example, if you are trying to use the HuggingFace providers out-of-the-box, you may encounter the error:

ValidationError: 1 validation error for HuggingFaceHub
__root__
  Could not import huggingface_hub python package. Please it install it with `pip install huggingface_hub`. (type=value_error)

The solution is quite simple; just run pip install huggingface_hub and restart your IPython kernel.

We list openai as an explicit dependency of Jupyter AI, so the OpenAI model providers should work out-of-the-box.

Loading the extension

You need to load the extension manually every time you start a new IPython session.

%load_ext jupyter_ai

Authentication

Before using models from a model provider, you usually must supply authentication credentials via an environment variable. For example, the openai and openai-chat providers expect the OPENAI_API_KEY environment variable to be set. This can be specified when starting the server:

OPENAI_API_KEY=<token> jupyter lab ...

Or can be declared in your shell profile file, which is ~/.bash_profile for Bash users and ~/.zprofile for Zsh users: (note that you will have to restart your shell sessions for this to take effect)

export OPENAI_API_KEY=<token>

Or can be done imperatively inside an IPython notebook cell:

%env OPENAI_API_KEY=<token>

Specifying model ID

A model ID can take three different forms:

1. Global model ID -- a model ID that is prefixed by its model provider ID.

• Model providers are sourced from langchain.llms, and their ID is defined under the _llm_type property of each provider class. Currently we support:
  • ai21, anthropic, cohere, huggingface_hub, huggingface_endpoint, openai, openai-chat, sagemaker_endpoint
• Example: openai-chat:gpt-3.5-turbo

2. Local model ID -- a model ID without a provider ID prefixed.

• Example: gpt-3.5-turbo

3. Model alias -- a shorthand alias to a model ID.

• Example: chatgpt is a model alias for gpt-3.5-turbo

When provided with a local model ID, we make a best-effort attempt to infer the provider ID. We do a search through all of the providers we support, and if one provider lists support for the model ID in question, we assume that to be the provider.

• For example: only openai-chat supports gpt-3.5-turbo, so we assume the provider to be openai-chat.

Registry providers

Some providers, like are unique in the sense that:

1. They do not have a static number of models they support. New models are being uploaded to HF Hub daily.
2. The model ID isn't necessarily a unique identifier for a model, but may be a URL to a model hosted on an endpoint.
3. The model ID has a different syntax for each of these providers

We call these providers "registry providers" (since they mimic the behavior of a package registry, like NPM/PyPi). They are defined internally with the models schema field set to [*] like so:

"sagemaker_endpoint": {
    "models": ["*"],
    "model_id_key": "endpoint_name",
},

Right now, we have 3 registry providers: huggingface_hub, huggingface_endpoint, sagemaker_endpoint. For each of these, the model ID has the syntax:

• huggingface_hub: the HF repo ID, e.g. gpt2 or google/flan-t5-xxl
• huggingface_endpoint: the URL to the HF endpoint, e.g. https://foo.us-west-2.aws.endpoints.huggingface.cloud/bar
• sagemaker_endpoint: the endpoint name. Your region and authentication credentials should already be specified via boto3.

We never assume a local model ID to be scoped to a registry provider. That means to use one of these registry providers, you must specify it explicitly by providing a global model ID:

%%ai huggingface_hub:flan-t5-xxl
Who won the World Cup in 2002?

This is because given just the local model ID google/flan-t5-xxl, it's ambiguous as to whether this should use the huggingface_hub or the sagemaker_endpoint provider.

When given just an HTTP/HTTPS URL, we assume it to be a HuggingFace endpoint, as that is the only supported registry provider that accepts raw URLs as a model ID.

Specifying format

Format currently includes the following: html, markdown, md, math, json. The format argument determines what IPython display is used to render output on the notebook. Currently we default to markdown, for which md is a short-hand alias.

Prompt interpolation via IPython

By default, any tokens wrapped in curly braces {} are replaced with the value of any variable of the same name in the current IPython scope. For example, running:

In [1]: poet = "Walt Whitman"
In [2]: %ai chatgpt Write a poem in the style of {poet}.

The prompt body is then interpolated to:

Write a poem in the style of Walt Whitman.

before being processed further and sent to the model.

Conversational usage

By default, we maintain a transcript for the openai-chat model provider, so the model is forwarded your previous exchanges for context. This means you can ask it questions about previous exchanges.

%%ai chatgpt
Who won the World Cup in 2002?

...

%%ai chatgpt
Tell me more about the coach of that team.

If you wish to discard the transcript for your current IPython session, use the --reset flag:

%ai chatgpt --reset

Note that we do not currently prune the transcript if it grows excessively large; so you may have to do this manually if you have an extensive number of exchanges within an IPython session.

Known issues

• For some reason, passing -f/--format when using %ai as a line magic triggers an argument parsing UsageError at runtime:

• The kernel hangs if the API is rate-limiting the user. There should be an enforced and configurable timeout.

• Prompt templates should be inferred per-format. For example, if --format math or --format json are specified, then the prompt template should be something similar to "Perform exclusively the task requested of you. Do not include a description.\n{body}".